#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
课本图片转HTML工具

功能：
1. 识别课本图片中的文字内容
2. 分析文本布局（标题、段落、列表等）
3. 生成保持原样式的HTML文件

使用方法：
1. 将课本图片放在同一目录下
2. 运行本程序，指定图片文件名
3. 查看生成的HTML文件
"""

import os
import sys
import cv2
import numpy as np
import pytesseract
from PIL import Image
import easyocr
import re
import argparse
from pathlib import Path

# Windows环境下需要设置Tesseract路径
# 请根据实际安装路径调整
pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

class BookPageProcessor:
    """课本页面处理器"""
    
    def __init__(self, use_easyocr=True):
        """初始化处理器
        
        Args:
            use_easyocr: 是否使用EasyOCR（更好的中文识别，但速度较慢）
        """
        self.use_easyocr = use_easyocr
        if use_easyocr:
            # 初始化EasyOCR阅读器（支持中英文）
            self.reader = easyocr.Reader(['ch_sim', 'en'])
    
    def load_image(self, image_path):
        """加载并预处理图像
        
        Args:
            image_path: 图像文件路径
        
        Returns:
            preprocessed_img: 预处理后的图像
            original_img: 原始图像
        """
        # 读取图像
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"无法读取图像: {image_path}")
        
        # 保存原始图像
        original_img = image.copy()
        
        # 预处理图像以提高OCR质量
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 自适应阈值处理，提高对比度
        processed = cv2.adaptiveThreshold(
            gray, 
            255,
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY,
            11, 2
        )
        
        # 降噪
        processed = cv2.medianBlur(processed, 3)
        
        return processed, original_img
    
    def detect_layout(self, image):
        """检测图像布局结构
        
        Args:
            image: 预处理后的图像
            
        Returns:
            regions: 识别到的区域列表，每个元素为(x, y, w, h, region_type)
        """
        # 复制图像以进行处理
        img = image.copy()
        
        # 查找轮廓
        contours, hierarchy = cv2.findContours(
            255 - img, 
            cv2.RETR_EXTERNAL,
            cv2.CHAIN_APPROX_SIMPLE
        )
        
        # 过滤并分析轮廓
        regions = []
        img_h, img_w = img.shape[:2]
        min_area = img_h * img_w * 0.005  # 最小区域大小（图像面积的0.5%）
        
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            area = w * h
            
            # 过滤太小的区域
            if area < min_area:
                continue
                
            # 判断区域类型（标题、段落、图片等）
            # 这里使用简单启发式规则，实际应用中可能需要更复杂的分析
            region_type = 'paragraph'  # 默认为段落
            
            # 如果区域宽度接近图像宽度且高度较小，可能是标题
            if w > img_w * 0.7 and h < img_h * 0.1:
                region_type = 'heading'
            
            # 如果区域宽高比例较大，可能是图表
            if w / h > 3 or h / w > 3:
                region_type = 'figure'
            
            regions.append((x, y, w, h, region_type))
        
        # 按从上到下的顺序排序区域
        regions.sort(key=lambda r: r[1])
        
        return regions
    
    def recognize_text(self, img, region=None):
        """识别图像或区域中的文本
        
        Args:
            img: 图像
            region: 可选，要识别的区域(x, y, w, h)
            
        Returns:
            text: 识别的文本
            confidence: 置信度
        """
        if region:
            x, y, w, h = region
            roi = img[y:y+h, x:x+w]
        else:
            roi = img
        
        if self.use_easyocr:
            # 使用EasyOCR识别文本
            results = self.reader.readtext(roi)
            
            # 提取文本和置信度
            texts = []
            confidence = 0
            
            for detection in results:
                bbox, text, conf = detection
                texts.append(text)
                confidence += conf
            
            if results:
                confidence /= len(results)  # 计算平均置信度
            else:
                confidence = 0
                
            return ' '.join(texts), confidence
        else:
            # 使用Tesseract OCR
            # 对于中文识别，使用chi_sim语言包
            text = pytesseract.image_to_string(roi, lang='chi_sim+eng')
            return text, 1.0  # Tesseract没有直接提供置信度
    
    def process_image(self, image_path):
        """处理图像并提取文本和布局
        
        Args:
            image_path: 图像文件路径
            
        Returns:
            content: 包含文本和布局信息的字典
        """
        # 加载并预处理图像
        processed_img, original_img = self.load_image(image_path)
        
        # 检测布局
        regions = self.detect_layout(processed_img)
        
        # 收集内容
        content = {
            'title': '',
            'elements': []
        }
        
        # 处理每个区域
        for i, (x, y, w, h, region_type) in enumerate(regions):
            # 从原始图像中提取区域
            region_img = original_img[y:y+h, x:x+w]
            
            # 识别文本
            text, confidence = self.recognize_text(original_img, (x, y, w, h))
            
            # 如果是第一个标题区域，将其设为内容标题
            if i == 0 and region_type == 'heading':
                content['title'] = text.strip()
            
            # 添加到元素列表
            element = {
                'type': region_type,
                'text': text.strip(),
                'confidence': confidence,
                'position': (x, y, w, h),
                'image': region_img  # 保存区域图像以备后用
            }
            content['elements'].append(element)
        
        return content
    
    def generate_html(self, content, output_path, include_images=True):
        """根据内容生成HTML文件
        
        Args:
            content: 包含文本和布局信息的字典
            output_path: 输出HTML文件路径
            include_images: 是否包含原图区域的图像
            
        Returns:
            html_path: 生成的HTML文件路径
        """
        # 创建输出目录（如果不存在）
        output_dir = os.path.dirname(output_path)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        # 如果需要包含图像，创建图像目录
        images_dir = None
        if include_images:
            images_dir = os.path.join(output_dir, 'region_images')
            if not os.path.exists(images_dir):
                os.makedirs(images_dir)
        
        # 构建HTML内容
        html = []
        html.append('<!DOCTYPE html>')
        html.append('<html lang="zh-CN">')
        html.append('<head>')
        html.append('<meta charset="UTF-8">')
        html.append('<meta name="viewport" content="width=device-width, initial-scale=1.0">')
        html.append(f'<title>{content["title"] or "课本页面"}</title>')
        html.append('<style>')
        html.append('''
            body {
                font-family: "Microsoft YaHei", SimSun, Arial, sans-serif;
                line-height: 1.6;
                margin: 0;
                padding: 20px;
                background-color: #f8f8f8;
            }
            .page-container {
                max-width: 800px;
                margin: 0 auto;
                background-color: white;
                padding: 30px;
                box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
            }
            .region {
                position: relative;
                margin-bottom: 20px;
                border-radius: 5px;
                overflow: hidden;
            }
            .heading {
                font-size: 24px;
                font-weight: bold;
                margin-bottom: 20px;
                color: #333;
                text-align: center;
            }
            .paragraph {
                font-size: 16px;
                margin-bottom: 15px;
                text-indent: 2em;
                line-height: 1.8;
            }
            .figure {
                text-align: center;
            }
            .figure img {
                max-width: 100%;
                margin: 10px 0;
                border: 1px solid #eee;
            }
            .region-image {
                max-width: 100%;
                border: 1px solid #ddd;
                margin-top: 5px;
            }
            .confidence {
                font-size: 12px;
                color: #999;
                margin-top: 5px;
                text-align: right;
            }
        ''')
        html.append('</style>')
        html.append('</head>')
        html.append('<body>')
        
        html.append('<div class="page-container">')
        
        # 添加标题
        if content['title']:
            html.append(f'<h1 class="heading">{content["title"]}</h1>')
        
        # 添加内容元素
        for i, element in enumerate(content['elements']):
            element_type = element['type']
            text = element['text']
            confidence = element['confidence']
            
            # 跳过空文本
            if not text.strip():
                continue
                
            # 根据元素类型添加不同的HTML标签
            html.append(f'<div class="region {element_type}">')
            
            if element_type == 'heading':
                html.append(f'<h2>{text}</h2>')
            elif element_type == 'paragraph':
                # 分段处理
                paragraphs = text.split('\n\n')
                for para in paragraphs:
                    if para.strip():
                        html.append(f'<p>{para}</p>')
            elif element_type == 'figure':
                html.append(f'<p>{text}</p>')
                
            # 如果需要，添加原区域图像
            if include_images and images_dir:
                # 保存区域图像
                image_filename = f'region_{i}.jpg'
                image_path = os.path.join(images_dir, image_filename)
                cv2.imwrite(image_path, element['image'])
                
                # 添加图像到HTML
                rel_image_path = os.path.join('region_images', image_filename)
                html.append(f'<img class="region-image" src="{rel_image_path}" alt="{element_type} region" />')
            
            # 添加置信度信息（可选）
            html.append(f'<div class="confidence">识别置信度: {confidence:.2f}</div>')
            
            html.append('</div>')
        
        html.append('</div>')  # 关闭page-container
        
        html.append('</body>')
        html.append('</html>')
        
        # 写入HTML文件
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(html))
        
        return output_path

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='将课本图片转换为HTML')
    parser.add_argument('image_path', help='课本图片路径')
    parser.add_argument('--output', '-o', help='输出HTML文件路径', default='book_page.html')
    parser.add_argument('--no-images', action='store_true', help='不包含原区域图像')
    parser.add_argument('--tesseract', action='store_true', help='使用Tesseract OCR（代替EasyOCR）')
    
    args = parser.parse_args()
    
    # 创建处理器
    processor = BookPageProcessor(use_easyocr=not args.tesseract)
    
    try:
        print(f"处理图像: {args.image_path}")
        # 处理图像
        content = processor.process_image(args.image_path)
        
        # 生成HTML
        html_path = processor.generate_html(
            content, 
            args.output,
            not args.no_images
        )
        
        print(f"已生成HTML文件: {html_path}")
        print(f"识别到 {len(content['elements'])} 个区域")
        
        # 显示识别结果摘要
        if content['title']:
            print(f"页面标题: {content['title']}")
        
        print("\n识别内容摘要:")
        for i, element in enumerate(content['elements']):
            if i < 3:  # 只显示前三个元素
                print(f"- {element['type']}: {element['text'][:50]}...")
        
        if len(content['elements']) > 3:
            print(f"... 还有 {len(content['elements']) - 3} 个区域")
            
    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
        return 1
    
    return 0

if __name__ == "__main__":
    sys.exit(main())